package profiler;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;

public class AgentBootstrap implements ClassFileTransformer {

    private Filter filter;

    public static final String INSTRUMENTATION_PREFFIX = "instrumented%";
    
    public AgentBootstrap(String config) {
        filter = new DefaultFilter();
    }

    public static void premain(String agentArguments, Instrumentation instrumentation) {
        AgentBootstrap agent = new AgentBootstrap(agentArguments);
        instrumentation.addTransformer(agent);
        instrumentation.setNativeMethodPrefix(agent, INSTRUMENTATION_PREFFIX);
    }

    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
            byte[] classfileBuffer) throws IllegalClassFormatException {
        if (classBeingRedefined != null) {
            if (filter.instrumentClass(className, classBeingRedefined.getModifiers())) {
                ClassReader reader = new ClassReader(classfileBuffer);
                ClassWriter writer = new ClassWriter(reader, 0);
                Instrumentor i = new Instrumentor(writer, filter);
                reader.accept(i, 0);
                byte[] byteArray = writer.toByteArray();

                dump(byteArray, className);

                return byteArray;
            } else {
                return classfileBuffer;
            }
        } else {
            ClassReader reader = new ClassReader(classfileBuffer);
            ClassWriter writer = new ClassWriter(reader, 0);
            Instrumentor i = new Instrumentor(writer, filter);
            reader.accept(i, 0);
            byte[] byteArray = writer.toByteArray();
            if(i.isAccepted()) {
//                ClassReader constructed = new ClassReader(classfileBuffer);
//                TraceClassVisitor visitor = new TraceClassVisitor(new PrintWriter(System.out));
//                constructed.accept(visitor, TraceClassVisitor.getDefaultAttributes(), ClassReader.SKIP_DEBUG);
            }
            return byteArray;
        }
    }

    public static void dump(byte[] byteArray, String className) {
        try {
            File name = new File(className + ".class");
            name.getParentFile().mkdirs();
            System.out.println("writing to: " + name);
            FileOutputStream fout = new FileOutputStream(name);
            fout.write(byteArray);
            fout.flush();
            fout.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

}
